#include <iomanip>
#include "Core/Wrapper_LAPACKE.h"
#include "Core/TensorExpr.h"
#include "Embedded/Convection.h"
#include "Embedded/SourceStencil.impl.h"
#include "Embedded/EvaluateMonomials.h"
#include "../example/common/TestUtility.h"

template <int Dim, int pD, int _fitOrder>
PartialDerOfProd<Dim, pD, _fitOrder>::PartialDerOfProd(const EmGrid &aGrid,
                                                       const Tensor<Real, Dim + 1> &aVel)
    : emGrid(aGrid)
{
  // 0.
  struct _partialDerTransform : public PartialDer<Dim, pD, _fitOrder>::partialDerTransform {
    enum { fitOrder = _fitOrder };
  };
  // 1. Construct the stencils.
  SourceStencilParam<Dim> ssparam;
  ssparam.otype = SourceStencilParam<Dim>::OperatorType::Convection;
  ssparam.velField = &aVel;
  CPUTimer timer;
  timer.reset();
  SourceStencil<Dim> ss(emGrid, _partialDerTransform(), ssparam);
  std::cout << "Construction of convection px" << pD << " in " << Dim << "D : "
            << std::setprecision(5) << timer() << "s. " << std::endl << std::defaultfloat;
#ifndef NDEBUG
  if(get_dbglevel() >= 2) {
    ss.dump(tmpos);
    tmpos.flush();
  }
#endif
  // 2. Find the approximation.
  const auto &bdryPoints = emGrid.bdryPoints;
  const auto dx = emGrid.spacing();
  const int n = fitOrder - 1;
  const int nBasis = binom(n+Dim, Dim);
  auto DirichletModifier = [&](const Real &x, const auto &) { return x; };

  timer.reset();
  for(const auto &stencil : ss.vStencil) {
    const int nSources = stencil.locations.size();
    const auto destIdx = stencil.destIdx;
    const auto refLoc  = (destIdx + 0.5) * dx;
    // 2.1. calculate the sample matrix.
    Tensor<Real, 2> sampleMatrix(Vec<int, 2> {nBasis, nSources});
    sampleMatrix = 0.0;
    Real *where2Fill = sampleMatrix.data();
    for(std::size_t k = 0; k < stencil.locations.size(); ++k) {
      rVec p;
      if(stencil.isBdry[k] == 0) {
        iVec jdx = stencil.locations[k];
        p = (jdx + 0.5) * dx;
        EvaluateMonomials<Dim,n>::doEvaluate(identityOfMonomials(), DirichletModifier, where2Fill, p-refLoc, rVec(0.0), dx);
      } else {
        // use the nearest boudary point as the boundary condition
        p = bdryPoints(destIdx);
        EvaluateMonomials<Dim,n>::doEvaluate(identityOfMonomials(), DirichletModifier, where2Fill, p-refLoc, rVec(0.0), dx);
      }
    }
    assert((where2Fill - sampleMatrix.data()) == sampleMatrix.volume());
    // 2.2. calculate the RHS.
    Tensor<Real, 1> fpd(nBasis);
    fpd = 0.0;
    where2Fill = fpd.data();
    EvaluateMonomials<Dim, n>::doEvaluate(_partialDerTransform(), DirichletModifier, where2Fill, rVec(0.0), rVec(0.0), dx);
    Tensor<Real, 1> fv(nBasis);
    fv = 0.0;
    where2Fill = fv.data();
    EvaluateMonomials<Dim, n>::doEvaluate(identityOfMonomials(), DirichletModifier, where2Fill, rVec(0.0), rVec(0.0), dx);
    Tensor<Real, 2> sampleRHS(Vec<int,2> {nSources, nSources});
    sampleRHS = 0.0;
    // sampleRHS is the (symmetric) tensor product of fv and fpd
    for(int j = 0; j < nBasis; ++j)
      for(int i = 0; i < nBasis; ++i)
        sampleRHS(i, j) = fpd(i) * fv(j) + fpd(j) * fv(i);
    // 2.3. solve the linear system to obtain the coefficients.
    // stage - 1
    for(int j = 0; j < nSources; ++j)
      for(int i = 0; i < nBasis; ++i)
        sampleMatrix(i, j) /= stencil.weights[j];
    Tensor<Real, 2> scaledSampleMatrix = sampleMatrix.copy();
    auto info = LAPACKE_dgels(LAPACK_COL_MAJOR, 'N', nBasis/*m*/, nSources/*n*/, nBasis/*nhrs*/, sampleMatrix.data(), nBasis/*lda*/, sampleRHS.data(), nSources/*ldb*/);
    if(info != 0)
      throw std::runtime_error("DGELS fails. ");
    for(int j = 0; j < nBasis; ++j)
      for(int i = 0; i < nSources; ++i)
        sampleRHS(i, j) /= stencil.weights[i];
    // stage - 2
    for(int j = 0; j < nSources; ++j) // make a transpose
      for(int i = 0; i < j; ++i)
        std::swap(sampleRHS(i, j), sampleRHS(j, i));
    info = LAPACKE_dgels(LAPACK_COL_MAJOR, 'N', nBasis/*m*/, nSources/*n*/, nSources/*nrhs*/, scaledSampleMatrix.data(), nBasis/*lda*/, sampleRHS.data(), nSources/*ldb*/);
    if(info != 0)
      throw std::runtime_error("DGELS fails. ");
    for(int j = 0; j < nSources; ++j)
      for(int i = 0; i < nSources; ++i)
        sampleRHS(i, j) /= stencil.weights[i];
    // 2.4. Save the coefficients in sparse matrix format.
    std::vector<iVV> LL(nSources * nSources);
    int m = 0;
    for(int j = 0; j < nSources; ++j)
      for(int i = 0; i < nSources; ++i)
        LL[m++] = std::make_pair(stencil.locations[i], stencil.locations[j]);
    matOp.insertRow(destIdx, nSources * nSources, LL.data(), sampleRHS.data());
    dbgcout2 << sampleRHS << "\n";
  }
  std::cout << "Discretization of the convection px" << pD << " in " << Dim << "D : "
            << std::setprecision(5) << timer() << "s. " << std::endl << std::defaultfloat;
}

template <int Dim, int pD, int _fitOrder>
void PartialDerOfProd<Dim, pD, _fitOrder>::apply(const Tensor<Real, Dim> &aData0,
                                                 const Tensor<Real, Dim> &aData1,
                                                 Tensor<Real, Dim> &outData) const
{
  // 1. apply the standard difference scheme.
  Box<Dim> validBox = emGrid;
  auto dx = emGrid.spacing();
  if(_fitOrder == 5) {
    ddfor(validBox, [&](const iVec &idx) {
      auto uj = iVec::unit(pD);
      Real r2 = ((aData0(idx + uj) + aData0(idx)) * (aData1(idx + uj) + aData1(idx))
          - (aData0(idx - uj) + aData0(idx)) * (aData1(idx - uj) + aData1(idx))) / (4 * dx[pD]);
      Real r4 = ((aData0(idx + uj * 2) + aData0(idx)) * (aData1(idx + uj * 2) + aData1(idx))
          - (aData0(idx - uj * 2) + aData0(idx)) * (aData1(idx - uj * 2) + aData1(idx))) / (8 * dx[pD]);
      outData(idx) = r2 * (4.0 / 3) - r4 * (1.0 / 3);
    });
  } else if(_fitOrder == 7) {
    auto uj = iVec::unit(pD);
    ddfor(validBox, [&](const iVec &idx) {
      Real r2 = ((aData0(idx + uj) + aData0(idx)) * (aData1(idx + uj) + aData1(idx))
          - (aData0(idx - uj) + aData0(idx)) * (aData1(idx - uj) + aData1(idx))) / (4 * dx[pD]);
      Real r4 = ((aData0(idx + uj * 2) + aData0(idx)) * (aData1(idx + uj * 2) + aData1(idx))
          - (aData0(idx - uj * 2) + aData0(idx)) * (aData1(idx - uj * 2) + aData1(idx))) / (8 * dx[pD]);
      Real r6 = ((aData0(idx + uj * 3) + aData0(idx)) * (aData1(idx + uj * 3) + aData1(idx))
          - (aData0(idx - uj * 3) + aData0(idx)) * (aData1(idx - uj * 3) + aData1(idx))) / (12 * dx[pD]);
      outData(idx) = r2 * (3.0 / 2) + r4 * (-3.0 / 5) + r6 * (1.0 / 10);
    });
  } else if(_fitOrder == 9) {
    auto uj = iVec::unit(pD);
    ddfor(validBox, [&](const iVec &idx) {
      Real r2 = ((aData0(idx + uj) + aData0(idx)) * (aData1(idx + uj) + aData1(idx))
          - (aData0(idx - uj) + aData0(idx)) * (aData1(idx - uj) + aData1(idx))) / (4 * dx[pD]);
      Real r4 = ((aData0(idx + uj * 2) + aData0(idx)) * (aData1(idx + uj * 2) + aData1(idx))
          - (aData0(idx - uj * 2) + aData0(idx)) * (aData1(idx - uj * 2) + aData1(idx))) / (8 * dx[pD]);
      Real r6 = ((aData0(idx + uj * 3) + aData0(idx)) * (aData1(idx + uj * 3) + aData1(idx))
          - (aData0(idx - uj * 3) + aData0(idx)) * (aData1(idx - uj * 3) + aData1(idx))) / (12 * dx[pD]);
      Real r8 = ((aData0(idx + uj * 4) + aData0(idx)) * (aData1(idx + uj * 4) + aData1(idx))
          - (aData0(idx - uj * 4) + aData0(idx)) * (aData1(idx - uj * 4) + aData1(idx))) / (16 * dx[pD]);
      outData(idx) = r2 * (8.0 / 5) + r4 * (-4.0 / 5) + r6 * (8.0 / 35) + r8 * (-1.0 / 35);
    });
  }
  // 2. apply the correction.
  matOp.AXPY(1.0, [&](const iVV &idxidx) -> Real {
    return aData0(idxidx.first) * aData1(idxidx.second);
  }, [](const iVec &) { return 0.0; }, outData);
}

//===============================================================================
template <int Dim, int fitOrder>
Convection<Dim, fitOrder>::Convection(const EmGrid &aGrid, const Tensor<Real, Dim + 1> &aVel)
    : emGrid(aGrid)
{
  static_for<0, Dim>::execute([&](auto _i_, int) {
    constexpr int i = decltype(_i_)::value;
    pdpObjs[i] = new PartialDerOfProd<Dim, i, fitOrder>(emGrid, aVel);
    return 0;
  }, 0);
}

template <int Dim, int fitOrder>
Convection<Dim, fitOrder>::~Convection()
{
  static_for<0, Dim>::execute([&](auto _i_, int) {
    constexpr int i = decltype(_i_)::value;
    delete reinterpret_cast<PartialDerOfProd<Dim, Dim-1-i, fitOrder> *>(pdpObjs[Dim-1-i]);
    return 0;
  }, 0);
}

template <int Dim, int fitOrder>
void Convection<Dim, fitOrder>::apply(const Tensor<Real, Dim + 1> &aVel, Tensor<Real, Dim + 1> &aCvt) const
{
  Tensor<Real, Dim> vel[] { {aVel, 0},
                            {aVel, 1},
                            {aVel, 2} };
  Tensor<Real, Dim> cvt[] { {aCvt, 0},
                            {aCvt, 1},
                            {aCvt, 2} };
  Tensor<Real, Dim> temp(cvt[0].box());
  for(int d = 0; d < Dim; ++d) {
    cvt[d] = 0.0;
    static_for<0, Dim>::execute([&](auto _i_, int) {
      constexpr int i = decltype(_i_)::value;
      reinterpret_cast<PartialDerOfProd<Dim, i, fitOrder> *>(pdpObjs[i])->apply(vel[d], vel[i], temp);
      cvt[d] = cvt[d] + temp;
      return 0;
    }, 0);
  }
}

//======================================================================
#if DIM == 2
template class Convection<2, 5>;
template class Convection<2, 7>;
template class Convection<2, 9>;
#else
template class Convection<3, 5>;
template class Convection<3, 7>;
#endif // DIM == 2
